import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QLineEdit, QTextEdit, QFileDialog, QMessageBox
import sqlite3
import openpyxl
import pandas as pd
import os

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("套管打印文件作成工具")
        self.setGeometry(100, 100, 1200, 1000)

        self.layout = QVBoxLayout()

        self.label = QLabel("你可以在此添加用于分隔套管的内容:")
        self.layout.addWidget(self.label)

        self.textbox = QLineEdit()
        self.layout.addWidget(self.textbox)

        # 添加清空数据库按钮
        self.clear_button = QPushButton("清空数据库")
        self.clear_button.clicked.connect(self.clear_database)
        self.layout.addWidget(self.clear_button)

        # 添加保存为Excel按钮
        self.save_button = QPushButton("保存为Excel")
        self.save_button.clicked.connect(self.save_to_excel)
        self.layout.addWidget(self.save_button)

        # 添加选择Excel文件按钮
        self.select_file_button = QPushButton("选择Excel文件")
        self.select_file_button.clicked.connect(self.load_excel_data)
        self.layout.addWidget(self.select_file_button)

        # 添加数据库显示区
        self.database_display = QTextEdit()
        self.layout.addWidget(self.database_display)

        self.setLayout(self.layout)

        # 连接到数据库
        self.conn = sqlite3.connect('database.db')
        self.c = self.conn.cursor()

        # 创建表
        self.create_table()

        # 初始显示数据库内容
        self.update_display()

    def create_table(self):
        self.c.execute("""
            CREATE TABLE IF NOT EXISTS DATA (
                GRAPHIC TEXT,
                `GROUP` TEXT,
                ID TEXT,
                TIME TEXT,
                MEMBER TEXT
            )
        """)
        self.conn.commit()

    def insert_data(self, graphic, group, id_value, time, member):
        self.c.execute("""
            INSERT INTO DATA (GRAPHIC, `GROUP`, ID, TIME, MEMBER)
            VALUES (?, ?, ?, ?, ?)
        """, (graphic, group, id_value, time, member))
        self.conn.commit()
        self.update_display()


    def update_display(self):
        self.c.execute("SELECT * FROM DATA")
        result = self.c.fetchall()
        self.database_display.clear()
        for row in result:
            self.database_display.append(str(row))

    def save_to_excel(self):
        try:
            self.c.execute("SELECT * FROM DATA")
            data = self.c.fetchall()
            df = pd.DataFrame(data, columns=['GRAPHIC', 'GROUP', 'ID', 'TIME', 'MEMBER'])

            file_name, _ = QFileDialog.getSaveFileName(self, "保存Excel文件", "", "Excel文件 (*.xlsx)")

            if file_name:
                df.to_excel(file_name, index=False)
                print("数据已成功保存为Excel文件。")
        except Exception as e:
            print("Error:", str(e))

    def clear_database(self):
        try:
            self.c.execute("DELETE FROM DATA")
            self.conn.commit()
            self.update_display()
            print("数据库已清空。")
        except Exception as e:
            print("Error:", str(e))

    def load_excel_data(self):
        file_names, _ = QFileDialog.getOpenFileNames(self, "选择Excel文件", "", "Excel文件 (*.xlsx *.xls)")
        for file_name in file_names:
            try:
                excel_name = os.path.splitext(os.path.basename(file_name))[0]
                workbook = openpyxl.load_workbook(file_name)
                worksheet = workbook.active  # 使用第一个工作表

                data = []
                for row in worksheet.iter_rows(min_row=2, max_col=53, values_only=True):  # BA列是第53列
                    graphic = row[1]  # B列
                    group = row[5]  # F列
                    id_value = row[6]  # G列
                    time = row[52]  # BA列
                    data.append((graphic, group, id_value, time, excel_name))

                self.c.executemany("""
                    INSERT INTO DATA (GRAPHIC, `GROUP`, ID, TIME, MEMBER)
                    VALUES (?, ?, ?, ?, ?)
                """, data)
                self.conn.commit()

                self.update_display()
                print("Excel数据已成功插入数据库。")
            except Exception as e:
                print("Error:", str(e))

if __name__ == "__main__":
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())

